// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/lib/tests/SliceTestBase.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>

namespace kumo::pollux::functions::prestosql {
namespace {

class SliceTest : public test::SliceTestBase {
 protected:
  static const vector_size_t kVectorSize{1000};
  static void SetUpTestCase() {
    SliceTestBase::SetUpTestCase();
    registerAllScalarFunctions();
  }
};

TEST_F(SliceTest, basic) {
  basicTestCases();
}

TEST_F(SliceTest, constantInputArray) {
  auto array_vector = make_array_vector<int64_t>({
      {1, 2, 3, 4, 5, 6, 7},
      {1, 2, 7},
      {1, 2, 3, 5, 6, 7},
  });

  // Happy test 1.
  {
    auto expectedArrayVector =
        make_array_vector<int64_t>({{2, 3}, {2, 7}, {2, 3}});
    testSlice("slice(C0, 2, 2)", {array_vector}, expectedArrayVector);
  }

  // Happy test 2.
  {
    auto expectedArrayVector = make_array_vector<int64_t>({{3}, {7}, {3}});

    testSlice("slice(C0, 3, 1)", {array_vector}, expectedArrayVector);
  }

  // Allow negative start index.
  {
    auto expectedArrayVector =
        make_array_vector<int64_t>({{6, 7}, {2, 7}, {6, 7}});

    testSlice("slice(C0, -2, 2)", {array_vector}, expectedArrayVector);
  }

  // Allow length to extend beyond boundary and be a BIGINT.
  {
    auto expectedArrayVector = make_array_vector<int64_t>(
        {{1, 2, 3, 4, 5, 6, 7}, {1, 2, 7}, {1, 2, 3, 5, 6, 7}});
    testSlice("slice(C0, 1, 2147483648)", {array_vector}, expectedArrayVector);
  }

  // Throw invalid argument when start index = 0.
  {
    auto expectedArrayVector = make_array_vector<int64_t>(
        {{1, 2, 3, 4, 5, 6, 7}, {1, 2, 7}, {1, 2, 3, 5, 6, 7}});
    POLLUX_ASSERT_THROW(
        testSlice("slice(C0, 0, 1)", {array_vector}, expectedArrayVector),
        "SQL array indices start at 1");
  }
}

TEST_F(SliceTest, variableInputArray) {
  {
    auto startsVector = make_flat_vector<int64_t>(
        kVectorSize, [](vector_size_t row) { return 1 + row % 7; });
    auto lengthsVector = make_flat_vector<int64_t>(
        kVectorSize, [](vector_size_t /*row*/) { return 1; });
    auto sizeAt = [](vector_size_t /*row*/) { return 7; };
    auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
      return 1 + idx;
    };
    auto array_vector = make_array_vector<int64_t>(kVectorSize, sizeAt, value_at);

    auto expectedSizeAt = [](vector_size_t /*row*/) { return 1; };
    auto expectedValueAt = [](vector_size_t row, vector_size_t /*idx*/) {
      return 1 + row % 7;
    };
    auto expectedArrayVector =
        make_array_vector<int64_t>(kVectorSize, expectedSizeAt, expectedValueAt);
    testSlice(
        "slice(C0, C1, C2)",
        {array_vector, startsVector, lengthsVector},
        expectedArrayVector);
  }

  {
    auto startsVector = make_flat_vector<int64_t>(
        kVectorSize, [](vector_size_t row) { return 1 + row % 7; });
    auto lengthsVector = make_flat_vector<int64_t>(
        kVectorSize, [](vector_size_t /*row*/) { return 2; });
    auto sizeAt = [](vector_size_t /*row*/) { return 7; };
    auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
      return 1 + idx;
    };
    auto array_vector = make_array_vector<int64_t>(kVectorSize, sizeAt, value_at);

    auto expectedSizeAt = [](vector_size_t row) {
      return (row + 1) % 7 == 0 ? 1 : 2;
    };
    auto expectedValueAt = [](vector_size_t row, vector_size_t idx) {
      return 1 + row % 7 + idx;
    };
    auto expectedArrayVector =
        make_array_vector<int64_t>(kVectorSize, expectedSizeAt, expectedValueAt);
    testSlice(
        "slice(C0, C1, C2)",
        {array_vector, startsVector, lengthsVector},
        expectedArrayVector);
  }

  {
    auto startsVector = make_flat_vector<int64_t>(
        kVectorSize, [](vector_size_t row) { return 1 + row % 7; });
    auto lengthsVector = make_flat_vector<int64_t>(
        kVectorSize, [](vector_size_t /*row*/) { return 7; });
    auto sizeAt = [](vector_size_t /*row*/) { return 7; };
    auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
      return 1 + idx;
    };
    auto array_vector = make_array_vector<int64_t>(kVectorSize, sizeAt, value_at);

    auto expectedSizeAt = [](vector_size_t row) { return 7 - row % 7; };
    auto expectedValueAt = [](vector_size_t row, vector_size_t idx) {
      return 1 + row % 7 + idx;
    };
    auto expectedArrayVector =
        make_array_vector<int64_t>(kVectorSize, expectedSizeAt, expectedValueAt);
    testSlice(
        "slice(C0, C1, C2)",
        {array_vector, startsVector, lengthsVector},
        expectedArrayVector);
  }
}

TEST_F(SliceTest, varcharVariableInput) {
  auto startsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t row) { return 1 + row % 7; });
  auto lengthsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 7; });

  auto sizeAt = [](vector_size_t /*row*/) { return 7; };
  auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
    return StringView::makeInline(melon::to<std::string>(idx + 1));
  };
  auto array_vector = make_array_vector<StringView>(kVectorSize, sizeAt, value_at);

  auto expectedSizeAt = [](vector_size_t row) { return 7 - row % 7; };
  auto expectedValueAt = [](vector_size_t row, vector_size_t idx) {
    return StringView::makeInline(melon::to<std::string>(1 + row % 7 + idx));
  };
  auto expectedArrayVector =
      make_array_vector<StringView>(kVectorSize, expectedSizeAt, expectedValueAt);

  testSlice(
      "slice(C0, C1, C2)",
      {array_vector, startsVector, lengthsVector},
      expectedArrayVector);
}

TEST_F(SliceTest, allIndicesGreaterThanArraySize) {
  auto startsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 8; });
  auto lengthsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 7; });
  auto sizeAt = [](vector_size_t /*row*/) { return 7; };
  auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
    return 1 + idx;
  };
  auto array_vector = make_array_vector<int64_t>(kVectorSize, sizeAt, value_at);

  auto expectedSizeAt = [](vector_size_t /*row*/) { return 0; };
  // Not going to be used, created to satisfy function signature.
  auto expectedValueAt = [](vector_size_t row) { return 1 + row % 7; };
  auto expectedArrayVector =
      make_array_vector<int64_t>(kVectorSize, expectedSizeAt, expectedValueAt);
  testSlice(
      "slice(C0, C1, C2)",
      {array_vector, startsVector, lengthsVector},
      expectedArrayVector);
}

TEST_F(SliceTest, errorStatesArray) {
  auto startsVector =
      make_flat_vector<int64_t>(kVectorSize, [](vector_size_t row) {
        // Zero out index -> should throw an error.
        if (row == 40) {
          return 0;
        }
        return 1 + row % 7;
      });
  auto lengthsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 7; });

  auto sizeAt = [](vector_size_t /*row*/) { return 7; };
  auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
    return StringView::makeInline(melon::to<std::string>(idx + 1));
  };
  auto array_vector = make_array_vector<StringView>(kVectorSize, sizeAt, value_at);

  auto expectedSizeAt = [](vector_size_t row) { return 7 - row % 7; };
  auto expectedValueAt = [](vector_size_t row, vector_size_t idx) {
    return StringView::makeInline(melon::to<std::string>(1 + row % 7 + idx));
  };
  auto expectedArrayVector =
      make_array_vector<StringView>(kVectorSize, expectedSizeAt, expectedValueAt);

  POLLUX_ASSERT_THROW(
      testSlice(
          "slice(C0, C1, C2)",
          {array_vector, startsVector, lengthsVector},
          expectedArrayVector),
      "SQL array indices start at 1");
}

TEST_F(SliceTest, zeroSliceLength) {
  auto startsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t row) { return 1 + row % 7; });
  auto lengthsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 0; });

  auto sizeAt = [](vector_size_t /*row*/) { return 7; };
  auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
    return StringView::makeInline(melon::to<std::string>(idx + 1));
  };
  auto array_vector = make_array_vector<StringView>(kVectorSize, sizeAt, value_at);

  auto expectedSizeAt = [](vector_size_t /*row*/) { return 0; };
  auto expectedValueAt = [](vector_size_t /*row*/, vector_size_t /*idx*/) {
    return StringView::makeInline(melon::to<std::string>());
  };
  auto expectedArrayVector =
      make_array_vector<StringView>(kVectorSize, expectedSizeAt, expectedValueAt);

  testSlice(
      "slice(C0, C1, C2)",
      {array_vector, startsVector, lengthsVector},
      expectedArrayVector);
}

TEST_F(SliceTest, negativeSliceLength) {
  auto startsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t row) { return 1 + row % 7; });
  auto lengthsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return -1; });
  auto sizeAt = [](vector_size_t /*row*/) { return 7; };
  auto value_at = [](vector_size_t /*row*/, vector_size_t idx) {
    return 1 + idx;
  };
  auto array_vector = make_array_vector<int64_t>(kVectorSize, sizeAt, value_at);

  // Not actually used.
  auto expectedSizeAt = [](vector_size_t row) { return 7 - row % 7; };
  auto expectedValueAt = [](vector_size_t row, vector_size_t idx) {
    return 1 + row % 7 + idx;
  };
  auto expectedArrayVector =
      make_array_vector<int64_t>(kVectorSize, expectedSizeAt, expectedValueAt);
  POLLUX_ASSERT_THROW(
      testSlice(
          "slice(C0, C1, C2)",
          {array_vector, startsVector, lengthsVector},
          expectedArrayVector),
      "The value of length argument of slice() function should not be negative");
}

TEST_F(SliceTest, constantArrayNonConstantLength) {
  // Tests constant arrays and non-constant starts and lengths. Ensure they
  // don't create overlapping ranges in the output ArrayVector.
  auto startsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 2; });
  auto lengthsVector = make_flat_vector<int64_t>(
      kVectorSize, [](vector_size_t /*row*/) { return 2; });
  auto array_vector = make_constant_array<int64_t>(kVectorSize, {99, 100, 101});

  auto expectedSizeAt = [](vector_size_t /*row*/) { return 2; };
  auto expectedValueAt = [](vector_size_t /*row*/, vector_size_t idx) {
    return idx == 0 ? 100 : 101;
  };
  auto expectedArrayVector =
      make_array_vector<int64_t>(kVectorSize, expectedSizeAt, expectedValueAt);
  testSlice(
      "slice(C0, C1, C2)",
      {array_vector, startsVector, lengthsVector},
      expectedArrayVector);
}

} // namespace
} // namespace kumo::pollux::functions::prestosql
